package org.ala.tiktools.rpc.plugin.load_balance;

import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;

import org.ala.tiktools.logs.LoggerFactory;
import org.ala.tiktools.rpc.config.RpcConfigInfo;
import org.ala.tiktools.rpc.plugin.node.NodeKey;
import org.slf4j.Logger;
import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;


/**
 * @author ala
 * @date 2025-04-10 14:07
 */
public abstract class AbstractDynamicLoadBalance extends AbstractLoadBalance {
    /**
     * port配置项
     */
    public static final String METADATA_PORT = "rpc-port";


    static Logger log = LoggerFactory.getLogger("rpc", "loadBalance");


    /**
     *  group中的节点列表
     */
    protected Map<String, Set<NodeKey>> nodeGroup = new ConcurrentHashMap<>();


    @Override
    public final void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        super.setApplicationContext(applicationContext);
        init(context);

        //  开启定时任务
        long delay = RpcConfigInfo.INSTANCE.getReNodeInfoInterval();
        schedule.schedule(new ScanNodeTask(this), delay, TimeUnit.MILLISECONDS);
    }
    protected void init(ApplicationContext context) {}

    /**
     * 从group中拉取节点信息
     */
    @Override
    protected void loadRpcNodeKeyByGroup(String group, Consumer<NodeKey> consumer) {
        Set<NodeKey> nodes = getNodes(group);
        if (nodes != null) {
            for (NodeKey node : nodes) {
                consumer.accept(node);
            }
        }
    }


    /**
     * 刷节点信息
     */
    protected abstract Set<NodeKey> loadNodeKeyFromRemote(String group);


    /**
     *  上次拉取的节点列表
     */
    protected Map<String, Set<NodeKey>> lastGroup = new HashMap<>();
    /**
     *  从上次拉取的节点列表中获取节点
     */
    protected Set<NodeKey> getNodes(String group) {
        Set<NodeKey> nodes = lastGroup.get(group);
        if (nodes == null) {
            synchronized (group.intern()) {
                nodes = lastGroup.get(group);
                if (nodes == null) {
                    Set<NodeKey> nks = loadNodeKeyFromRemote(group);
                    if (nks != null) {
                        lastGroup.put(group, nks);
                        nodes = nks;
                    }
                }
            }
        }
        return nodes;
    }

    /**
     *  对比更新lastGroup
     */
    protected void comparisonAndUpdate(Set<NodeKey> newNks, String group) {
        Set<NodeKey> oldNks = lastGroup.get(group);
        if (oldNks != null) {
            //  newNks 比 oldNks 多出来的部分。这部分append
            for (NodeKey nk : newNks) {
                if (!oldNks.contains(nk)) {
                    log.info("[{} comparisonAndUpdate] 发现新增节点 node:{}", this.getClass().getSimpleName(), nk);
                    try {
                        appendNode(group, nk);
                    } catch (Exception e) {
                        log.error("[{} comparisonAndUpdate] 新增节点失败 node:{}", this.getClass().getSimpleName(), nk, e);
                    }
                }
            }
            //  oldNks 比 newNks 多出来的部分。这部分remove
            for (NodeKey nk : oldNks) {
                if (!newNks.contains(nk)) {
                    log.info("[{} comparisonAndUpdate] 发现摘除节点 node:{}", this.getClass().getSimpleName(), nk);
                    try {
                        removeNode(nk);
                    } catch (Exception e) {
                        log.error("[{} comparisonAndUpdate] kill节点失败 node:{}", this.getClass().getSimpleName(), nk, e);
                    }
                }
            }
        } else {
            for (NodeKey nk : newNks) {
                log.info("[{} comparisonAndUpdate] 发现新增节点 node:{}", this.getClass().getSimpleName(), nk);
                try {
                    appendNode(group, nk);
                } catch (Exception e) {
                    log.error("[{} comparisonAndUpdate] 新增节点失败 node:{}", this.getClass().getSimpleName(), nk, e);
                }
            }
        }

        lastGroup.put(group, newNks);
    }


    /**
     *  扫描节点定时任务
     */
    public static class ScanNodeTask implements Runnable {
        protected AbstractDynamicLoadBalance loadBalance;
        public ScanNodeTask(AbstractDynamicLoadBalance loadBalance) {
            this.loadBalance = loadBalance;
        }
        @Override
        public void run() {
            try {
                log.info("[{}} ScanNodeTask] 拉取节点信息", this.getClass().getSimpleName());
                //  从活跃列表中拿group
                Set<String> groups = loadBalance.activeNodes.keySet();

                for (String group : groups) {
                    //  更新节点列表
                    Set<NodeKey> nks = loadBalance.loadNodeKeyFromRemote(group);
                    if (nks == null) {continue;}
                    log.info("[{}} ScanNodeTask] 本次拉取到的节点 group:{} nodes:{}", this.getClass().getSimpleName(), group, nks);
                    //  和老的对比，多出来的append，少的remove
                    loadBalance.comparisonAndUpdate(nks, group);
                }
            } catch (Exception e) {
                log.error("[{}} ScanNodeTask] 拉取任务失败", this.getClass().getSimpleName(), e);
            } finally {
                schedule.schedule(this, RpcConfigInfo.INSTANCE.getReNodeInfoInterval(), TimeUnit.MILLISECONDS);
            }
        }
    }
}
